import sys
import os
from ConfigParser import SafeConfigParser
from socket import *
from threading import Thread, Event, Lock
import time
import paramiko
import logging
import StringIO

# IMPORTANT!! MAKE SURE LOGGING IS SET TO 0 or 2 WHEN RUNNING AS A SERVICE.
# SERVICE WILL HANG IF TRYING TO WRITE TO STDOUT.  (bbservice.py sets it to
# 2 prior to starting bumblebee as a service.)
LOGGING = 2 #0=no logging, 1=stdout, 2=log file, 3=stdout and log file
PRIVATE_RSA_KEY_SOURCE = 0 # 0=embedded in this file using function getPrivateRSAKey()
                           # 1=found in key file "EXEC_PATH\config\id_rsa"

def find_exec_path():
    if hasattr(sys,"frozen") and sys.frozen == "windows_exe":
        return os.path.dirname(os.path.abspath(sys.executable))
    else:
        return os.path.abspath(os.path.join(os.path.dirname(__file__)))

EXEC_PATH = find_exec_path()

def log(s):
    if LOGGING == 1 or LOGGING == 3:
        # log output to either stdout or file
        print '%s: %s' % (time.ctime(), s)
        sys.stdout.flush()
    if LOGGING >= 2:
        # log output to both.
        logging.info(s)

class PipeThread( Thread ):
    """This thread simply handles passing traffic from one socket to another.\n
    It listens on 'source' and sends what it hears to 'sink'."""
    global_halt_flag = Event() # if set, all threads of this type will kill themselves and close their connections.
    
    def __init__( self, source, sink, halt_flag_poll_interval=5):
        Thread.__init__( self )
        self.source = source # who it listens to
        self.source.settimeout(halt_flag_poll_interval) # timeout to allow the thread to check halt flags
        self.sink = sink # passes the data on to this guy.
        self.paired_thread = None # Another thread handles data coming back from 'sink'.
                                  # We keep a reference to it here so that we can signal it to
                                  # kill itself when this thread dies, since they operate
                                  # in pairs, and a thread has no other way to know when
                                  # its corresponding thread has died.
        self.halt_flag = Event()
        log( 'Creating new pipe thread  %s ( %s -> %s )' % \
            ( self, source.getpeername(), sink.getpeername() ))

    def run( self ):
        # keep listening for data to pass along as long as no halt flags are set.
        while not self.halt_flag.isSet() and not PipeThread.global_halt_flag.isSet():
            while not self.halt_flag.isSet() and not PipeThread.global_halt_flag.isSet():
                # handle 1024 bytes at a time, then come around again.
                try:
                    data = None
                    data = self.source.recv(1024)
                    break
                except:
                    pass

            try:
                if not data: break # die when the source has sent EOF
                self.sink.send(data)
            except:
                break

        log( '%s terminating' % self )
        self.source.close()
        self.sink.close()

        # kill corresponding thread, which otherwise doesn't know to die.
        if self.paired_thread:
            self.paired_thread.halt()
            self.paired_thread = None

    def halt(self):
        log("Halt flag set for pipe thread  %s" % self)
        self.halt_flag.set()

class Bumblebee(Thread):
    def __init__(self, config_path):
        Thread.__init__(self)

        # Load configuration settings
        config_parser = SafeConfigParser()
        config_parser.read(config_path)
        self.config = dict(config_parser.items("DEFAULT"))
                           
        #Initialize settings
        self.halt_flag = Event()

        #Setup logging
        log_dir = os.path.normpath(os.path.join(EXEC_PATH, self.config["log_dir"]))
        if not os.path.exists(log_dir):
            os.mkdir(log_dir)        
        log_path = os.path.normpath(os.path.join(log_dir, self.config["log_filename"]))
        logging.basicConfig(level=logging.INFO,
                            format='%(asctime)s %(levelname)s %(message)s',
                            filename=log_path,
                            filemode='a')

        #bumblebee settings
        self.pacetrac_server_addr = self.config["remote_server_addr"]
        self.pacetrac_server_port = int(self.config["remote_server_port"])
        self.pacetrac_forwarding_port = int(self.config["remote_server_forwarding_port"])
        self.bumblebee_addr = self.config["bumblebee_addr"]
        self.bumblebee_port = int(self.config["bumblebee_port"])
        self.target_port = int(self.config["target_port"])
        self.username = self.config["username"]
        self.seconds_before_retry = int(self.config["seconds_before_retry"])
        self.keepalive_interval = int(self.config["keepalive_interval"])
        self.halt_flag_polling_interval = int(self.config["halt_flag_polling_interval"])
        self.channel_acceptance_timeout = 10 #in seconds

        # the connection to pacetrac
        self.transport = None

        self.pipes = []

    def __getPrivateRSAKey(self):
        keydata = ""#insert your key data here
        s = StringIO.StringIO(keydata)
        return paramiko.RSAKey(file_obj = s)        

    def attemptConnection(self):
        """Attempts to connect to pacetrac server and setup reverse port forwarding.\n
        Returns true if successful, in which case self.transport is setup properly.\n
        Otherwise, returns false.\n"
        Uses RSA key to authenticate.  Set PRIVATE_RSA_KEY_SOURCE to set whether\n
        the key is embedded in this file (in getPrivateRSAKey()) or in a file.\n
        This file's path is expected to be EXEC_PATH\config\id_rsa"""
        try:
            log("Bumblebee attempting to connect to Remote Server at %s:%d" % (self.pacetrac_server_addr, self.pacetrac_server_port))

            self.transport = None

            # open a socket
            try:
                sock = socket(AF_INET, SOCK_STREAM)
                sock.connect((self.pacetrac_server_addr, self.pacetrac_server_port))
            except Exception, e:
                log("*** Connect failed: " + str(e))
                return False

            # try to negotiate SSH connection
            try:
                self.transport = paramiko.Transport(sock)
                self.transport.start_client()
            except paramiko.SSHException:
                log("*** SSH negotiation failed.")
                return False

            # load known hosts
            try:
                known_hosts = paramiko.util.load_host_keys(os.path.expanduser("~/.ssh/known_hosts"))
            except IOError:
                try:
                    known_hosts = paramiko.util.load_host_keys(os.path.expanduser("~/ssh/known_hosts"))
                except IOError:
                    log("*** Unable to open host keys file")
                    known_hosts = {}
            
            # check the host key against known hosts
            host_key = self.transport.get_remote_server_key()
            if not known_hosts.has_key(self.pacetrac_server_addr):
                log("*** WARNING: Unknown host key!  Continuing anyway...")
            elif not known_hosts[self.pacetrac_server_addr].has_key(host_key.get_name()):
                log("*** WARNING: Unknown host key!  Continuing anyway...")
            elif known_hosts[self.pacetrac_server_addr][host_key.get_name()] != host_key:
                log("*** WARNING: Host key has changed!  Dropping connection.")
                return False
            else:
                log("*** Host key OK.")

            if PRIVATE_RSA_KEY_SOURCE == 0:
                private_key = self.__getPrivateRSAKey()
            elif PRIVATE_RSA_KEY_SOURCE == 1:
                try:
                    # get key from file; does not expect it to be password-protected.
                    path = os.path.join(EXEC_PATH, "config", "id_rsa")
                    private_key = paramiko.RSAKey.from_private_key_file(path)
                except:
                    log("*** Could not access private RSA key file.")
                    return False

            # try to authenticate
            try:
                self.transport.auth_publickey(self.username, private_key)
                log("*** Successful authentication with Remote Server.")
            except:
                log("*** Could not authenticate private key.")
                return False

            if not self.transport.is_authenticated():
                log("*** Could not authenticate private key.")
                return False

            # Request port forwarding            
            log("Bumblebee is requesting port forwarding (port %d) from Remote Server..." % self.pacetrac_forwarding_port)
            try:
                port_assigned = self.transport.request_port_forward("%s:%s" % (self.bumblebee_addr, self.bumblebee_port), self.pacetrac_forwarding_port)
                log("Remote Server server is now forwarding its own port %d to Bumblebee." % port_assigned)
            except Exception, e:
                log("Bumblebee's port forwarding request denied: " + str(e))
                return False

            # set keepalive pulse
            log("*** Setting keepalive pulse to fire every %d seconds" % self.keepalive_interval)
            self.transport.set_keepalive(self.keepalive_interval)

            return True
        except Exception, e:
            log("Exception in attemptConnection(): " + str(e))
            return False

    def waitToRetryConnection(self):
        """Stalls for a preset number of seconds.  Checks the halt flag once a second, though,\n
        in case we are trying to shut down Bumblebee."""
        log("Error tunneling to Remote Server.  Waiting %d seconds before retry..." % self.seconds_before_retry)
        i = 0
        while not self.halt_flag.isSet() and i < self.seconds_before_retry:
            time.sleep(1)
            i += 1

    def clearDeadPipes(self):
        try:
            for pipe in self.pipes:
                if not pipe.isAlive():
                    self.pipes.remove(pipe)
        except Exception, e:
            log("Exception in clearDeadPipes(): " + str(e))

    def run(self):
        log("Starting bumblebee thread.")
        try:
            # try to connect until successful.
            while not self.halt_flag.isSet():
                try:
                    if self.attemptConnection():
                        log("Bumblebee redirecting: localhost:%s -> %s:%s" % (self.bumblebee_port, "localhost", self.target_port))
                        log("Bumblebee is listening...")                
                        while not self.halt_flag.isSet():                            
                            try:
                                new_channel = None
                                new_channel = self.transport.accept(self.channel_acceptance_timeout)
                                if new_channel:
                                    log("Bumblebee opened new channel from Remote Server")
                                    # link the incoming channel to another channel handling the incoming data
                                    # It is either an SSH server on another port, or Primus
                                    fwd = socket(AF_INET, SOCK_STREAM)
                                    fwd.connect(("localhost", self.target_port))
                                    p = PipeThread(new_channel, fwd, self.halt_flag_polling_interval)
                                    q = PipeThread(fwd, new_channel, self.halt_flag_polling_interval)
                                    # pair the threads so that if one dies, it can kill the other.
                                    p.paired_thread = q
                                    q.paired_thread = p
                                    self.pipes.append(p)
                                    self.pipes.append(q)
                                    log("%d pipes owned by bumblebee." % len(self.pipes))
                                    p.start()
                                    q.start()                            
                                    log("Bumblebee is listening...")
                                else:
                                    if not self.transport.is_active():                                        
                                        # if our transport has been disconnected,
                                        # start attempting to reconnect
                                        break                                    

                            except:
                                if not self.transport.is_active():
                                    # if an exception occured here that broke our
                                    # transport connection, start trying to
                                    # reconnect
                                    raise
                            
                            self.clearDeadPipes()
                    else:
                        if self.transport: self.transport.close()
                        self.waitToRetryConnection()
                except Exception, e:
                    log("Remote Server connection error: " + str(e))
                    if self.transport: self.transport.close()
                    self.waitToRetryConnection()
                
        except Exception, e:
            log("Bumblebee stopped listening: " + str(e))
            # shut down and release pipes, transport

        # we're halting, so make sure that all pipes are killed off
        # The only case we'll need this is if an exception occurred above.
        PipeThread.global_halt_flag.set()
        log("Waiting for PipeThreads to halt...")
        for pipe in self.pipes:
            log("waiting for pipe %s" % pipe)
            pipe.join()

        log("Bumblebee halted.")

    def halt(self):
        """Sets halt flags in both Bumblebee (to make it stop listening for\n
        connection requests) and all PipeThreads that are forwarding traffic.\n
        All threads will stop in a matter of seconds."""
        self.halt_flag.set()
        log("Bumblebee halt flag set")
        PipeThread.global_halt_flag.set()
        log("Global halt flag set for PipeThread.")

def main(argv=None):
    """A function to allow running bumblebee from the command line\n
    for testing.  Once it has started, type q<Enter> to stop."""
    if argv is None:
        argv = sys.argv

    config_path = os.path.join(EXEC_PATH, "config", "bumblebee.ini")
    bb = Bumblebee(config_path)

    bb.start()

    while True:
        s = raw_input("Type q<Enter> to quit.\n")
        if s == "q":
            bb.halt()
            break
        
    bb.join()
    logging.shutdown()
        
if __name__ == "__main__":
    sys.exit(main())  
